from LeetTool import TreeNode
from LeetTool import build_TreeNode


class Solution:
    def maxValue(self, root: TreeNode, k: int) -> int:
        return max(self.dfs(root, k))

    def dfs(self, node, k):
        if not node:
            return [0] + [float("-inf")] * k

        dp_left = self.dfs(node.left, k)
        dp_right = self.dfs(node.right, k)

        dp = [float("-inf")] * (k + 1)
        dp[0] = max(dp_left) + max(dp_right)

        for n in range(1, k + 1):
            for i in range(n):
                j = (n - 1) - i
                dp[n] = max(dp[n], node.val + dp_left[i] + dp_right[j])

        # print(node.val, ":", dp_left, dp_right, "->", dp)

        return dp


if __name__ == "__main__":
    print(Solution().maxValue(root=build_TreeNode([5, 2, 3, 4]), k=2))  # 12
    print(Solution().maxValue(root=build_TreeNode([4, 1, 3, 9, None, None, 2]), k=2))  # 16
